package metric

import (
	"strings"
	"sync"

	"gddgame.cc/galaxy/binocle"
	"github.com/toolkits/nux"
)

var (
	diskStatsMap = make(map[string][2]*nux.DiskStats)
	dsLock       = new(sync.RWMutex)
)

func UpdateDiskStats() error {
	dsList, err := nux.ListDiskStats()
	if err != nil {
		return err
	}

	dsLock.Lock()
	defer dsLock.Unlock()
	for i := 0; i < len(dsList); i++ {
		device := dsList[i].Device
		diskStatsMap[device] = [2]*nux.DiskStats{dsList[i], diskStatsMap[device][0]}
	}
	return nil
}

func IOReadRequests(arr [2]*nux.DiskStats) uint64 {
	return arr[0].ReadRequests - arr[1].ReadRequests
}

func IOReadMerged(arr [2]*nux.DiskStats) uint64 {
	return arr[0].ReadMerged - arr[1].ReadMerged
}

func IOReadSectors(arr [2]*nux.DiskStats) uint64 {
	return arr[0].ReadSectors - arr[1].ReadSectors
}

func IOMsecRead(arr [2]*nux.DiskStats) uint64 {
	return arr[0].MsecRead - arr[1].MsecRead
}

func IOWriteRequests(arr [2]*nux.DiskStats) uint64 {
	return arr[0].WriteRequests - arr[1].WriteRequests
}

func IOWriteMerged(arr [2]*nux.DiskStats) uint64 {
	return arr[0].WriteMerged - arr[1].WriteMerged
}

func IOWriteSectors(arr [2]*nux.DiskStats) uint64 {
	return arr[0].WriteSectors - arr[1].WriteSectors
}

func IOMsecWrite(arr [2]*nux.DiskStats) uint64 {
	return arr[0].MsecWrite - arr[1].MsecWrite
}

func IOMsecTotal(arr [2]*nux.DiskStats) uint64 {
	return arr[0].MsecTotal - arr[1].MsecTotal
}

func IOMsecWeightedTotal(arr [2]*nux.DiskStats) uint64 {
	return arr[0].MsecWeightedTotal - arr[1].MsecWeightedTotal
}

func TS(arr [2]*nux.DiskStats) uint64 {
	return uint64(arr[0].TS.Sub(arr[1].TS).Nanoseconds() / 1000000)
}

func IODelta(device string, f func([2]*nux.DiskStats) uint64) uint64 {
	val, ok := diskStatsMap[device]
	if !ok {
		return 0
	}

	if val[1] == nil {
		return 0
	}
	return f(val)
}

func ShouldHandleDevice(device string) bool {
	normal := len(device) == 3 && (strings.HasPrefix(device, "sd") || strings.HasPrefix(device, "vd"))
	aws := len(device) >= 4 && strings.HasPrefix(device, "xvd")
	flash := len(device) >= 4 && (strings.HasPrefix(device, "fio") || strings.HasPrefix(device, "nvme"))
	return normal || aws || flash
}

func Disk(agent binocle.Service) error {
	if err := UpdateDiskStats(); err != nil {
		return err
	}
	agent.AddMetrics().Collect(func() map[string]float32 {
		_ = UpdateDiskStats()
		dsLock.RLock()
		defer dsLock.RUnlock()
		var rio uint64 = 0
		var wio uint64 = 0
		var delta_rsec uint64 = 0
		var delta_wsec uint64 = 0
		var ruse uint64 = 0
		var wuse uint64 = 0
		var use uint64 = 0
		var duration uint64 = 0
		var weighted uint64 = 0

		for device := range diskStatsMap {
			if !ShouldHandleDevice(device) {
				continue
			}

			rio += IODelta(device, IOReadRequests)
			wio += IODelta(device, IOWriteRequests)
			delta_rsec += IODelta(device, IOReadSectors)
			delta_wsec += IODelta(device, IOWriteSectors)
			ruse += IODelta(device, IOMsecRead)
			wuse += IODelta(device, IOMsecWrite)
			use += IODelta(device, IOMsecTotal)
			duration += IODelta(device, TS)
			weighted += IODelta(device, IOMsecWeightedTotal)
		}

		n_io := rio + wio
		avgrq_sz := float32(0)
		await := float32(0)
		svctm := float32(0)
		if n_io != 0 {
			avgrq_sz = float32(delta_rsec+delta_wsec) / float32(n_io)
			await = float32(ruse+wuse) / float32(n_io)
			svctm = float32(use) / float32(n_io)
		}
		tmp := float32(use) * 100.0 / float32(duration)
		if tmp > 100.0 {
			tmp = 100.0
		}
		return map[string]float32{
			"disk.io.read_bytes":  float32(delta_rsec) * 512.0,
			"disk.io.write_bytes": float32(delta_wsec) * 512.0,
			"disk.io.avgrq_sz":    avgrq_sz,
			"disk.io.avgqu_sz":    float32(weighted) / 1000.0,
			"disk.io.await":       await * 512.0,
			"disk.io.svctm":       svctm * 512.0,
			"disk.io.util":        tmp,
		}
	})
	return nil
}
